import sys
sys.path.append("..")
sys.path.append("../..")
from RumdetecFramework.AdverRumorFramework import *
from Dataloader.twitterloader import *
from Dataloader.dataloader_utils import *
from SentModel.Sent2Vec import *
from PropModel.SeqPropagation import GRUModel
import torch
import torch.nn as nn

def pretrained_LSTMRD(pretrained_file):
    lvec = LSTMVec("../../saved/bert_embedding.pkl", "../../../bert_en/", 768, 768)
    prop = GRUModel(768, 256, 1, 0.2)
    cls = nn.Linear(256, 2)
    ch = torch.load(pretrained_file)
    lvec.load_state_dict(ch['sent2vec'])
    prop.load_state_dict(ch['prop_model'])
    cls.load_state_dict(ch['rdm_cls'])
    AdverRD = PropAdverRumorDetection(lvec, prop, cls, topic_label_num=5)
    return AdverRD

def obtain_general_set(tr_prefix, dev_prefix, te_prefix):
    tr_set = TwitterSet()
    tr_set.load_data_fast(data_prefix=tr_prefix, min_len=5)
    dev_set = TwitterSet()
    dev_set.load_data_fast(data_prefix=dev_prefix)
    te_set = TwitterSet()
    te_set.load_data_fast(data_prefix=te_prefix)
    return tr_set, dev_set, te_set

for i in range(5):
    tr, dev, te = obtain_general_set("../../data/twitter_tr%d"%i, "../../data/twitter_dev%d"%i, "../../data/twitter_te%d"%i)
    tr.filter_short_seq(min_len=5)
    tr.ResortSample(6)
    print("%s : (dev event)/(test event)/(train event) = %3d/%3d/%3d" % (te.data[te.data_ID[0]]['event'], len(dev), len(te), len(tr)))
    print("\n\n===========SubRDM Train===========\n\n")
    model = pretrained_LSTMRD("../../saved/LSTMRD_%s.pkl"%te.data[te.data_ID[0]]['event'])
    model.AdverTrainIters(tr, dev, te, max_epochs=100,
                    log_dir="../../logs/", log_suffix="_AdverLSTMRD",
                    model_file="P_AdverLSTMRD1_%s.pkl"% (te.data[te.data_ID[0]]['event']))